import numpy as np

from manimlib import *


class TestGlsl(Scene):
    def construct(self):
        arrow = Arrow(LEFT * 2, RIGHT * 2)
        arrow.set_color_by_code(f"""
            vec3 blue = vec3{tuple(hex_to_rgb(BLUE))};
            vec3 red = vec3{tuple(hex_to_rgb(RED))};
            color.rgb = mix(blue, red, (point.x + 1.5) / 3);
        """)
        print(tuple(hex_to_rgb(BLUE)))
        self.add(arrow)


class ColorMapScene(Scene):
    def construct(self):
        radius = 3.
        sq = Square(side_length=2 * radius, fill_opacity=1)
        sq.set_color_by_code(f"""
            vec3 white = vec3(1.0);

            float pi = 3.14159;
            float tau = pi * 2.0;
            float theta_1 = atan(point.y / point.x);
            float theta_2 = atan(tan(theta_1 + 1.0 / 3.0 * tau));
            float theta_3 = atan(tan(theta_2 + 1.0 / 3.0 * tau));
        
            vec3 edge = vec3(theta_1 / pi, theta_2 / pi, theta_3 / pi);
            color.rgb = mix(white, edge, length(point / 3.5));
        """)
        self.add(sq)


class GlCode(Scene):
    def construct(self):
        code = Code("""def square_array(arr):
    return list(map(lambda x: x ** 2, arr))""",
                    font="Monaco_fix")
        self.add(code)


class TestSetColor(Scene):
    def construct(self):
        sq1 = DotCloud()
        sq1.to_grid(300, 300, 1).scale(1.5)
        sq1.set_radius(0.02)
        sq1.set_color(WHITE)
        sq1.set_color_by_code('''
        color.a = 1 / exp(length(point));
        ''')
        # sq1.make_3d()
        self.add(sq1)


class TestGradientColor(Scene):
    def construct(self):
        squ = RegularPolygon(n=4, fill_opacity=1, stroke_width=1).scale(4)
        squ.set_color(['#ff00ff', '#0000ff', '#0000ff', '#000000', '#000000', '#ff0000', '#ff0000', '#ffff00'])
        self.add(squ)


class LightPoints(Scene):
    def construct(self):
        text = Text("Hello World").scale(5)
        bg = BackgroundRectangle(text, buff=1)
        px = ParametricSurface(
            lambda u, v: [u, v, 0],
            u_range=(bg.get_left()[0], bg.get_right()[0]),
            v_range=(bg.get_bottom()[1], bg.get_top()[1]), color=WHITE)

        code = """color.a = 0.0;"""

        for point in text.get_all_points()[::3]:
            code += f"color.a += 0.003/(length((point-vec3{tuple(point)})*2))- 0.0005;"

        px.set_color_by_code(code)
        self.add(px)
        self.add(text)


class GlowDot(ParametricSurface):
    CONFIG = {
        'radius': 0.25,
        'resolution': (50, 50),
        'color': WHITE,
        'location': ORIGIN
    }

    def __init__(self, **kwargs):
        digest_config(self, kwargs)
        super().__init__(self.uv_func, u_range=(0, self.radius), v_range=(0, TAU))
        self.init_color()

    def uv_func(self, u, v):
        return np.array([u * np.cos(v), u * np.sin(v), 0])

    def init_color(self):
        self.set_color_by_code(f'''
            color.a=1/exp(length((point-vec3{tuple(self.location)})*{self.radius * 2}));
        ''')


class TestGlowDot(Scene):
    def construct(self):
        p = ORIGIN
        d = TrueDot(radius=1, color=WHITE)
        d.set_color_by_code(f'''
            color.a=1/exp(length((point-vec3{tuple(p)})*{d.get_radius() * 5}))- 0.004;
        ''')
        self.add(d)
        self.play(d.animate.shift(UP))


from manim_sandbox.utils.functions.debugTeX import PointIndex


class TestCamera(Scene):
    def construct(self):
        vector = Arrow().scale(5)
        pnts = PointIndex(vector, scale_factor=0.1)
        self.add(vector, pnts)


class Blur(ParametricSurface):
    def __init__(self, radius=2, k=1.5, color="#ffffdd", **kwargs):
        super(Blur, self).__init__(lambda u, v: v * np.array([
            np.cos(u), np.sin(u), 0
        ]), u_range=[0, TAU], v_range=[0, radius], **kwargs)

        rgba_func = lambda p: [*hex_to_rgb(color), 1 / np.exp(get_norm(p * k)) - 0.07]
        points = self.get_points()
        rgba_arrays = [rgba_func(point) for point in points[:len(points) // 3]]
        self.set_rgba_array(rgba_arrays)


class TestBlurDot(Scene):
    def construct(self):
        blur = Blur()
        self.add(blur)
        self.play(blur.shift, UP)
        self.wait()
        self.play(blur.shift, DOWN)
        self.wait()

